import pandas as pd
import math
import matplotlib.pyplot as plt
from pylab import mpl   # 从pylab导入子模块mpl
mpl.rcParams['font.sans-serif'] = ['KaiTi']    # 以楷体的字体显示中文
mpl.rcParams['axes.unicode_minus'] = False  # 解决保存图像是负号'-'显示为方块的问题


class Visualization:

    def get_yield(self, recorddf, cash):
        """
        函数作用：得到策略的累计收益率的DataFrame，由资产价值获得
        :param recorddf: dataFrame，index是日期，columns是当日所持资产市值
        :param cash: 初始金额
        :return: dataFrame, 策略累计收益率
        """
        recorddf["策略的累计收益率"] = recorddf["资产价值"].map(lambda x: x / cash - 1)
        yielddf = pd.DataFrame(recorddf["策略的累计收益率"])
        return yielddf

    def sum_return_single(self, SZ_data, starttime, endtime, title, return_BaseOn):
        """
        函数作用：绘制出某股票 在所选时间段 累计收益率 的曲线图。
                当日累计收益率可以通过收盘价来计算，当日累计收益率sum_return=当日收盘价/初始日期的收盘价-1；
               当日累计收益率还可以通过当日所持资产的市值来计算，当日累计收益率sum_return=当日所持资产的市值/初始日期所持资产的市值-1

        :param SZ_data: dataFrame，index是日期，columns是收盘价、当日所持资产市值等指标
        :param starttime: str, 设定的实施策略的开始时间，默认是"2019-07-01"
        :param endtime: str, 设定的实施策略的结束时间，默认是"2019-12-31"
        :param title: str, 输出的曲线图的标题
        :param return_BaseOn: 用来计算累计收益率的指标
                            对基准收益(上证指数、沪深300）的累计收益率进行计算和可视化的时候，累计收益率是基于“收盘价”计算的
                            对自己策略的累计收益率进行计算和可视化的时候，累计收益率是基于“所持资产当日市值”计算的
        :return: dataFrame, 基准的累计收益率
        """

        SZ_close_data = SZ_data.loc[starttime:endtime, return_BaseOn]  # 截取所选时间段的所需指标return_BaseOn信息，得到的是序列
        SZ_close_data = pd.DataFrame(SZ_close_data)  # 将序列转换为数据框
        start_date_closeprice = SZ_close_data.loc[starttime, return_BaseOn]  # 获取起始日期的收盘价 # 索引到这一行这一列的元素
        SZ_close_data['所选股票当日累计收益率'] = SZ_close_data[return_BaseOn].map(
            lambda x: (x / start_date_closeprice) - 1)  # 在数据框的最右侧追加一列累计收益率，累计收益率由“当日收盘价/起始日期收盘价-1”获得
        SZ_SumReturn = SZ_close_data.loc[starttime:endtime, "所选股票当日累计收益率"]  # 截取截取所选时间段的"当日累计收益率"指标
        SZ_SumReturn_df = pd.DataFrame(SZ_SumReturn)  # 将序列转换为数据框格式
        plt.figure(figsize=(9, 6))
        plt.plot(SZ_SumReturn_df)  # 画该数据框的折线图
        plt.title(title)  # 将该折线图的标题命名为"title"
        plt.xticks(fontsize=10, rotation=0)
        plt.yticks(fontsize=10)
        plt.xlabel(u'选取的时间段', fontsize=13, rotation=0)  # x轴的标题为"选取的时间段"
        plt.ylabel(u'累计收益率', fontsize=13, rotation=90)  # y轴的标题为"累计收益率"
        plt.grid(True)
        plt.show()
        return SZ_SumReturn_df


    def factor_visual(self, Stock_data, starttime, endtime, title, factor):
        """
        函数作用：绘制出某股票 在所选时间段 具体某个因子数值变化 的时间序列曲线图
        :param Stock_data: dataFrame，index是日期，columns是因子名称
        :param starttime: str, 设定的实施策略的开始时间，默认是"2019-07-01"
        :param endtime: str, 设定的实施策略的结束时间，默认是"2019-12-31"
        :param title: str, 输出的曲线图的标题
        :param factor: str, 要绘制的因子，是Stock_data的某个columns
        """

        Stock_factor = Stock_data.loc[starttime:endtime, factor]  # 截取所选时间段的factor因子的信息，得到的是序列
        Stock_factor = pd.DataFrame(Stock_factor)  # 将序列转换为数据框
        plt.figure(figsize=(9, 6))
        plt.plot(Stock_factor)  # 绘制该数据框的折线图
        plt.title(title)  # 将该折线图的标题命名为"title"
        plt.xticks(fontsize=10, rotation=0)
        plt.yticks(fontsize=10)
        plt.xlabel(u'选取的时间段', fontsize=13, rotation=0)  # x轴的标题为"选取的时间段"
        plt.ylabel(factor, fontsize=13, rotation=90)  # y轴的标题为"factor"
        plt.grid(True)
        plt.show()

